from bank_info import get_format_info
import pandas as pd

# 从CSV文件读取测试数据
df = pd.read_csv("data/测试句子.csv")

# 获取数据行数，处理所有句子
N = df.shape[0]

# 定义需要提取的信息字段
keys = [
    "银行",  # 银行名称
    "类型",  # 业务类型
    "金额",  # 金额数值
    "开始时间",  # 开始日期
    "结束时间",  # 结束日期
    "存放天数",  # 存放期限
    "补贴单价",  # 单位补贴价格
    "补贴类型",  # 补贴计算方式
    "总补贴金额",  # 总补贴金额
    "备注",  # 其他备注信息
]

# 为DataFrame添加空的结果列
for key in keys:
    df[key] = ""

# 逐行处理每个句子，提取银行信息
for i in range(N):
    txt = df.loc[i, "内容"]  # 获取原始文本内容

    # 调用信息提取函数
    result, formatted_text, res_info_str = get_format_info(txt)

    # 将提取结果填入对应列
    for key in keys:
        df.loc[i, key] = str(result.get(key, ""))

# 保存处理结果到新的CSV文件
df.to_csv("data/测试句子_结果_gpt41.csv", index=False)
